{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "411d50b8",
   "metadata": {
    "colab_type": "text",
    "id": "EgiF12Hf1Dhs"
   },
   "source": [
    "This notebook provides examples to go along with the [textbook](http://manipulation.csail.mit.edu/pose.html).  I recommend having both windows open, side-by-side!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "207638a3",
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "eeMrMI0-1Dhu",
    "lines_to_end_of_cell_marker": 2
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from pydrake.all import (\n",
    "    AbstractValue,\n",
    "    Concatenate,\n",
    "    DiagramBuilder,\n",
    "    LeafSystem,\n",
    "    PiecewisePolynomial,\n",
    "    PiecewisePose,\n",
    "    PointCloud,\n",
    "    RigidTransform,\n",
    "    RollPitchYaw,\n",
    "    Simulator,\n",
    "    StartMeshcat,\n",
    ")\n",
    "\n",
    "from manipulation import FindResource, running_as_notebook\n",
    "from manipulation.icp import IterativeClosestPoint\n",
    "from manipulation.meshcat_utils import AddMeshcatTriad\n",
    "from manipulation.mustard_depth_camera_example import MustardPointCloud\n",
    "from manipulation.pick import (\n",
    "    MakeGripperCommandTrajectory,\n",
    "    MakeGripperFrames,\n",
    "    MakeGripperPoseTrajectory,\n",
    ")\n",
    "from manipulation.scenarios import AddIiwaDifferentialIK\n",
    "from manipulation.station import (\n",
    "    AddPointClouds,\n",
    "    MakeHardwareStation,\n",
    "    load_scenario,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8dc1671c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Start the visualizer.\n",
    "meshcat = StartMeshcat()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a3dc18fc",
   "metadata": {},
   "source": [
    "# Putting it all together\n",
    "\n",
    "In the code above, we worked with a point cloud using functions.  To assemble this into a full-stack manipulation system, we need to specify the timing semantics of when those functions are called.  That's precisely what Drake's systems framework provides.  I've introduced two systems below:\n",
    "- `MustardIterativeClosestPoint` system that takes the camera inputs and outputs the pose estimate using ICP, and\n",
    "- `PickAndPlaceTrajectory` system that takes this pose estimate (and the state of the robot), computes the trajectory, and stores that trajectory in its Context so that it can output the instantaneous command.\n",
    "\n",
    "We don't use a `TrajectorySource` here, because the trajectory is not known when we first build the Diagram... the information we need to plan the trajectory requires reading sensors at runtime."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31c41ca9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Takes 3 point clouds (in world coordinates) as input, and outputs and estimated pose for the mustard bottle.\n",
    "class MustardIterativeClosestPoint(LeafSystem):\n",
    "    def __init__(self):\n",
    "        LeafSystem.__init__(self)\n",
    "        model_point_cloud = AbstractValue.Make(PointCloud(0))\n",
    "        self.DeclareAbstractInputPort(\"cloud0\", model_point_cloud)\n",
    "        self.DeclareAbstractInputPort(\"cloud1\", model_point_cloud)\n",
    "        self.DeclareAbstractInputPort(\"cloud2\", model_point_cloud)\n",
    "\n",
    "        self.DeclareAbstractOutputPort(\n",
    "            \"X_WO\",\n",
    "            lambda: AbstractValue.Make(RigidTransform()),\n",
    "            self.EstimatePose,\n",
    "        )\n",
    "\n",
    "        self.mustard = MustardPointCloud()\n",
    "        meshcat.SetObject(\"icp_scene\", self.mustard)\n",
    "\n",
    "    def EstimatePose(self, context, output):\n",
    "        pcd = []\n",
    "        for i in range(3):\n",
    "            cloud = self.get_input_port(i).Eval(context)\n",
    "            pcd.append(\n",
    "                cloud.Crop(\n",
    "                    lower_xyz=[0.4, -0.2, 0.001], upper_xyz=[0.6, 0.3, 0.3]\n",
    "                )\n",
    "            )\n",
    "        merged_pcd = Concatenate(pcd)\n",
    "        down_sampled_pcd = merged_pcd.VoxelizedDownSample(voxel_size=0.005)\n",
    "        meshcat.SetObject(\n",
    "            \"icp_observations\", down_sampled_pcd, point_size=0.001\n",
    "        )\n",
    "\n",
    "        X_WOhat, chat = IterativeClosestPoint(\n",
    "            self.mustard.xyzs(),\n",
    "            down_sampled_pcd.xyzs(),\n",
    "            meshcat=meshcat,\n",
    "            meshcat_scene_path=\"icp_scene\",\n",
    "        )\n",
    "\n",
    "        output.set_value(X_WOhat)\n",
    "\n",
    "\n",
    "class PickAndPlaceTrajectory(LeafSystem):\n",
    "    def __init__(self, plant):\n",
    "        LeafSystem.__init__(self)\n",
    "        self._gripper_body_index = plant.GetBodyByName(\"body\").index()\n",
    "        self.DeclareAbstractInputPort(\n",
    "            \"body_poses\", AbstractValue.Make([RigidTransform()])\n",
    "        )\n",
    "        self.DeclareAbstractInputPort(\n",
    "            \"X_WO\", AbstractValue.Make(RigidTransform())\n",
    "        )\n",
    "\n",
    "        self.DeclareInitializationUnrestrictedUpdateEvent(self.Plan)\n",
    "        self._traj_X_G_index = self.DeclareAbstractState(\n",
    "            AbstractValue.Make(PiecewisePose())\n",
    "        )\n",
    "        self._traj_wsg_index = self.DeclareAbstractState(\n",
    "            AbstractValue.Make(PiecewisePolynomial())\n",
    "        )\n",
    "\n",
    "        self.DeclareAbstractOutputPort(\n",
    "            \"X_WG\",\n",
    "            lambda: AbstractValue.Make(RigidTransform()),\n",
    "            self.CalcGripperPose,\n",
    "        )\n",
    "        self.DeclareVectorOutputPort(\"wsg_position\", 1, self.CalcWsgPosition)\n",
    "\n",
    "    def Plan(self, context, state):\n",
    "        X_G = {\n",
    "            \"initial\": self.get_input_port(0).Eval(context)[\n",
    "                int(self._gripper_body_index)\n",
    "            ]\n",
    "        }\n",
    "        X_O = {\n",
    "            \"initial\": self.get_input_port(1).Eval(context),\n",
    "            \"goal\": RigidTransform([0, -0.6, 0]),\n",
    "        }\n",
    "        X_GgraspO = RigidTransform(RollPitchYaw(np.pi / 2, 0, 0), [0, 0.22, 0])\n",
    "        X_OGgrasp = X_GgraspO.inverse()\n",
    "        X_G[\"pick\"] = X_O[\"initial\"] @ X_OGgrasp\n",
    "        X_G[\"place\"] = X_O[\"goal\"] @ X_OGgrasp\n",
    "        X_G, times = MakeGripperFrames(X_G)\n",
    "        print(f\"Planned {times['postplace']} second trajectory.\")\n",
    "\n",
    "        if False:  # Useful for debugging\n",
    "            AddMeshcatTriad(meshcat, \"X_Oinitial\", X_PT=X_O[\"initial\"])\n",
    "            AddMeshcatTriad(meshcat, \"X_Gprepick\", X_PT=X_G[\"prepick\"])\n",
    "            AddMeshcatTriad(meshcat, \"X_Gpick\", X_PT=X_G[\"pick\"])\n",
    "            AddMeshcatTriad(meshcat, \"X_Gplace\", X_PT=X_G[\"place\"])\n",
    "\n",
    "        traj_X_G = MakeGripperPoseTrajectory(X_G, times)\n",
    "        traj_wsg_command = MakeGripperCommandTrajectory(times)\n",
    "\n",
    "        state.get_mutable_abstract_state(int(self._traj_X_G_index)).set_value(\n",
    "            traj_X_G\n",
    "        )\n",
    "        state.get_mutable_abstract_state(int(self._traj_wsg_index)).set_value(\n",
    "            traj_wsg_command\n",
    "        )\n",
    "\n",
    "    def start_time(self, context):\n",
    "        return (\n",
    "            context.get_abstract_state(int(self._traj_X_G_index))\n",
    "            .get_value()\n",
    "            .start_time()\n",
    "        )\n",
    "\n",
    "    def end_time(self, context):\n",
    "        return (\n",
    "            context.get_abstract_state(int(self._traj_X_G_index))\n",
    "            .get_value()\n",
    "            .end_time()\n",
    "        )\n",
    "\n",
    "    def CalcGripperPose(self, context, output):\n",
    "        # Evaluate the trajectory at the current time, and write it to the\n",
    "        # output port.\n",
    "        output.set_value(\n",
    "            context.get_abstract_state(int(self._traj_X_G_index))\n",
    "            .get_value()\n",
    "            .GetPose(context.get_time())\n",
    "        )\n",
    "\n",
    "    def CalcWsgPosition(self, context, output):\n",
    "        # Evaluate the trajectory at the current time, and write it to the\n",
    "        # output port.\n",
    "        output.SetFromVector(\n",
    "            context.get_abstract_state(int(self._traj_wsg_index))\n",
    "            .get_value()\n",
    "            .value(context.get_time())\n",
    "        )\n",
    "\n",
    "\n",
    "def icp_pick_and_place_demo():\n",
    "    builder = DiagramBuilder()\n",
    "\n",
    "    scenario = load_scenario(\n",
    "        filename=FindResource(\"models/clutter.scenarios.yaml\"),\n",
    "        scenario_name=\"Mustard\",\n",
    "    )\n",
    "    station = builder.AddSystem(MakeHardwareStation(scenario, meshcat=meshcat))\n",
    "    to_point_cloud = AddPointClouds(\n",
    "        scenario=scenario, station=station, builder=builder\n",
    "    )\n",
    "\n",
    "    icp = builder.AddSystem(MustardIterativeClosestPoint())\n",
    "    builder.Connect(\n",
    "        to_point_cloud[\"camera3\"].get_output_port(), icp.get_input_port(0)\n",
    "    )\n",
    "    builder.Connect(\n",
    "        to_point_cloud[\"camera4\"].get_output_port(), icp.get_input_port(1)\n",
    "    )\n",
    "    builder.Connect(\n",
    "        to_point_cloud[\"camera5\"].get_output_port(), icp.get_input_port(2)\n",
    "    )\n",
    "\n",
    "    plant = station.GetSubsystemByName(\"plant\")\n",
    "    plan = builder.AddSystem(PickAndPlaceTrajectory(plant))\n",
    "    builder.Connect(\n",
    "        station.GetOutputPort(\"body_poses\"), plan.GetInputPort(\"body_poses\")\n",
    "    )\n",
    "    builder.Connect(icp.GetOutputPort(\"X_WO\"), plan.GetInputPort(\"X_WO\"))\n",
    "\n",
    "    robot = station.GetSubsystemByName(\n",
    "        \"iiwa.controller\"\n",
    "    ).get_multibody_plant_for_control()\n",
    "\n",
    "    # Set up differential inverse kinematics.\n",
    "    diff_ik = AddIiwaDifferentialIK(builder, robot)\n",
    "    builder.Connect(\n",
    "        diff_ik.get_output_port(), station.GetInputPort(\"iiwa.position\")\n",
    "    )\n",
    "    builder.Connect(plan.GetOutputPort(\"X_WG\"), diff_ik.get_input_port(0))\n",
    "    builder.Connect(\n",
    "        station.GetOutputPort(\"iiwa.state_estimated\"),\n",
    "        diff_ik.GetInputPort(\"robot_state\"),\n",
    "    )\n",
    "\n",
    "    builder.Connect(\n",
    "        plan.GetOutputPort(\"wsg_position\"),\n",
    "        station.GetInputPort(\"wsg.position\"),\n",
    "    )\n",
    "\n",
    "    diagram = builder.Build()\n",
    "\n",
    "    simulator = Simulator(diagram)\n",
    "    context = simulator.get_context()\n",
    "\n",
    "    simulator.Initialize()\n",
    "    if False:  # draw the trajectory triads\n",
    "        X_G_traj = (\n",
    "            plan.GetMyContextFromRoot(context)\n",
    "            .get_abstract_state(0)\n",
    "            .get_value()\n",
    "        )\n",
    "        for t in np.linspace(X_G_traj.start_time(), X_G_traj.end_time(), 40):\n",
    "            AddMeshcatTriad(\n",
    "                meshcat,\n",
    "                f\"X_G/({t})\",\n",
    "                X_PT=X_G_traj.GetPose(t),\n",
    "                length=0.1,\n",
    "                radius=0.004,\n",
    "            )\n",
    "\n",
    "    if running_as_notebook:\n",
    "        meshcat.StartRecording(set_visualizations_while_recording=True)\n",
    "        simulator.AdvanceTo(plan.end_time(plan.GetMyContextFromRoot(context)))\n",
    "        meshcat.PublishRecording()\n",
    "    else:\n",
    "        simulator.AdvanceTo(0.1)\n",
    "\n",
    "\n",
    "icp_pick_and_place_demo()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.10 64-bit",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
